# encoding: utf-8

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from MySqlTables import Base_data


class MyDB(object):
    def __init__(self, user='root', password='123456', address='localhost:3306', db='local.MyDB', base=None):
        self.user = user
        self.password = password
        self.address = address
        self.db = db
        self.base = base
        self.engine = self.create_engine()
        self.session = self.create_session()

    def create_engine(self):
        return create_engine("mysql+pymysql://{}:{}@{}/{}?charset={}".format(self.user, self.password, self.address, self.db, 'utf8'))

    def create_session(self):
        Session = sessionmaker(bind=self.engine)
        return Session()

    def init_db(self):
        self.base.metadata.create_all(self.engine)

    def drop_db(self):
        self.base.metadata.drop_all(self.engine)

    def connect(self):
        return self.engine.connect()


class MySqliteDB(MyDB):
    def __init__(self, user=None, password=None, address=None, db='local.MyDB', base=Base_data):
        super(MySqliteDB, self).__init__(user, password, address, db, base)
        self.engine = self.create_engine()

    def create_engine(self):
        return create_engine('sqlite:///{}'.format(self.db))